"""
MEMESA-TOOLS
============
Model preparation, vertex enumeration and vertex analysis tools developed as
part of the MEMESA project.

    Copyright (C) 2009-2012 Brett G. Olivier

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>

Author: Brett G. Olivier
Contact email: bgoli@users.sourceforge.net

"""

import os, time
cDir = os.path.dirname(os.path.abspath(os.sys.argv[0]))
import scipy.stats
import matplotlib
import matplotlib.pyplot as P
HD5datF = None


# Set userdata here #

# yeast 5.32
# work_dir = 'yeast532'
# model_name = 'yeast_5.33_model'  # base name not xml file
# #HD5datF = os.path.join(cDir, work_dir, 'vertex', 'yeast_5.33_model.None.hdf5')

# big Ecoli
work_dir = 'iAF1260'
model_name = 'Ecoli_iAF1260_noox.glc' # base name not xml file

# small Ecoli
# work_dir = 'iJR904'
# model_name = 'Ecoli_iJR904.ac' # base name not xml file
# model_name = 'Ecoli_iJR904.glc' # base name not xml file
# model_name = 'Ecoli_iJR904.glc2' # base name not xml file
# model_name = 'Ecoli_iJR904.lcts' # base name not xml file
# model_name = 'Ecoli_iJR904.succ' # base name not xml file

"""
Analysis methods
"""
if HD5datF == None:
    HD5datF = os.path.join(cDir, work_dir, 'vertex', '%s.None.hdf5' % model_name)

ENABLE_CORRCOEFF_METHOD_1 = True  # uses numpy.corrcoef(VarVertArray, y=None, rowvar=False, bias=False)
ENABLE_PEARSON_METHOD_1 = False   # uses scipy.stats.pearsonr(x, y)
ENABLE_SPEARMAN_METHOD_1 = False  # uses scipy.stats.spearmanr(x, y)

# end userdata #
model_file = '%s.xml' % model_name
model_dir = os.path.join(cDir, work_dir, 'sbml')
work_dir = os.path.join(cDir, work_dir, 'vertex')
vertex_variable_fluxes = "%s.vertex_variable_fluxes.csv" % model_name
vertex_variable_fluxes_all = "%s.vertex_variable_fluxes_all.csv" % model_name

import numpy
import h5py

from pyscescbm.CBVersion import __DEBUG__, __version__
from pyscescbm import CBRead, CBWrite, CBTools
from pyscescbm import CBSolver as slv

if not os.path.exists(work_dir):
    os.mkdir(work_dir)

mod = CBRead.readSBML2FBA(model_file, work_dir=model_dir)
mod.id = model_name

print(model_name)

CBTools.addStoichToFBAModel(mod)
CBTools.processBiGGchemFormula(mod)
CBTools.processSBMLAnnotationNotes(mod, annotation_key='note')

def getVariableFluxInfo(vertex_variable_fluxes, vertex_variable_fluxes_all, work_dir=work_dir):
    """
    Read in the varaible flux information and process them in terms of indexes, enums etc.
    Returns: VARIABLE_FLUX_NAMES, VARIABLE_FLUX_NAMES_idx, VARIABLE_FLUX_NAMES_enum
    """
    if work_dir != None:
        vertex_variable_fluxes = os.path.join(work_dir, vertex_variable_fluxes)
        vertex_variable_fluxes_all = os.path.join(work_dir, vertex_variable_fluxes_all)

    # get variable vertex flux names
    VF = file(vertex_variable_fluxes,'r')
    VARIABLE_FLUX_NAMES = tuple(VF.read().strip().split(','))
    print(VARIABLE_FLUX_NAMES)
    print(len(VARIABLE_FLUX_NAMES))
    VF.close

    # get variable vertex ennumerated flux names
    VF = file(vertex_variable_fluxes_all,'r')
    VFA = []
    for l in VF:
        VFA.append(l.strip().split(','))

    VARIABLE_FLUX_NAMES_enum = []
    VARIABLE_FLUX_NAMES_idx = []
    for j in range(len(VFA[0])):
        if VFA[1][j] == 'True':
            VARIABLE_FLUX_NAMES_enum.append((j,VFA[0][j]))
            VARIABLE_FLUX_NAMES_idx.append(j)
    del VFA
    VARIABLE_FLUX_NAMES_enum = tuple(VARIABLE_FLUX_NAMES_enum)
    VARIABLE_FLUX_NAMES_idx = numpy.array(VARIABLE_FLUX_NAMES_idx)
    print(VARIABLE_FLUX_NAMES_enum)
    print(len(VARIABLE_FLUX_NAMES_enum))
    print('Variable vertex idx:', len(VARIABLE_FLUX_NAMES_idx))
    VF.close
    return VARIABLE_FLUX_NAMES, VARIABLE_FLUX_NAMES_idx, VARIABLE_FLUX_NAMES_enum

VARIABLE_FLUX_NAMES, VARIABLE_FLUX_NAMES_idx, VARIABLE_FLUX_NAMES_enum = getVariableFluxInfo(vertex_variable_fluxes, vertex_variable_fluxes_all, work_dir=work_dir)

# open HDF5 vertex file
HD5dat = h5py.File(HD5datF,'r')
print(len(HD5dat['data/vertices']), HD5dat['data/vertices'].shape)

VertCorrCoeff = None
if ENABLE_CORRCOEFF_METHOD_1:
    TIME_START = time.time()
    ##  VarVertArray = numpy.zeros((HD5dat['data/vertices'].shape[0],len(VARIABLE_FLUX_NAMES_idx)),'f')
    VarVertArray = numpy.ones((HD5dat['data/vertices'].shape[0],len(VARIABLE_FLUX_NAMES_idx)),'d')

    cntr = 0
    print('Start time:', time.strftime('%H:%M:%S'))
    for idx in xrange(HD5dat['data/vertices'].shape[0]):
        VarVertArray[idx] =  HD5dat['data/vertices'][idx].take(VARIABLE_FLUX_NAMES_idx)
        cntr += 1
        if cntr == 10001:
            print(idx, float(idx)/float(HD5dat['data/vertices'].shape[0])*100.0, ((time.time()-TIME_START)/60.0), ((time.time()-TIME_START)/idx*float(HD5dat['data/vertices'].shape[0]))/60.0)
            cntr = 1

    print('VarVertArrayLen', len(VarVertArray))

    print('VarVertArr', VarVertArray.shape)
    VertCorrCoeff = numpy.corrcoef(VarVertArray, y=None, rowvar=False, bias=False)
    print('VertCorrCoeff', VertCorrCoeff.shape)
    print('CORRCOEFF_METHOD_1 time:', float(time.time()-TIME_START)/60.0, 'minutes')

def outputCorrCoef(CCarr, method_name, flux_names, model_name, work_dir):
    """
    Writes out the correlation coefficient matrix as a csv file and PNG figure.

     - *CCarr* the correlation coefficient matrix
     - *method_name* the name of the correlation algorithm e.g. 'spearman'
     - *model_name* the name of the current model
     - *work_dir* the working or output directory

    """
    CBTools.exportLabelledArray2CSV(CCarr, flux_names, os.path.join(work_dir, model_name+'.%s' % method_name))
    m = P.matshow(CCarr)
    P.savefig(os.path.join(work_dir, model_name+'.%s.png' % method_name), dpi=300)
    P.close('all')

if ENABLE_CORRCOEFF_METHOD_1:
    outputCorrCoef(VertCorrCoeff, 'corrcoeff', VARIABLE_FLUX_NAMES, model_name=model_name, work_dir=work_dir)

PearsonCorrCoeff = None
SpearmanCorrCoeff = None
if ENABLE_PEARSON_METHOD_1 or ENABLE_SPEARMAN_METHOD_1:
    TIME_START = time.time()
    if ENABLE_SPEARMAN_METHOD_1:
        SpearmanCorrCoeff = numpy.zeros((len(VARIABLE_FLUX_NAMES_idx),len(VARIABLE_FLUX_NAMES_idx)),'d')
        print('Init SpearmanCorrCoeff', SpearmanCorrCoeff.shape)
    if ENABLE_PEARSON_METHOD_1:
        PearsonCorrCoeff = numpy.zeros((len(VARIABLE_FLUX_NAMES_idx),len(VARIABLE_FLUX_NAMES_idx)),'d')
        print('Init PearsonCorrCoeff', PearsonCorrCoeff.shape)

    cntr = 0
    print('Start time:', time.strftime('%H:%M:%S'))
    for r in range(len(VARIABLE_FLUX_NAMES_idx)):
        rcol= HD5dat['data/vertices'][:,VARIABLE_FLUX_NAMES_idx[r]]
        for c in range(r, len(VARIABLE_FLUX_NAMES_idx)):
            if r == c:
                if ENABLE_SPEARMAN_METHOD_1:
                    SpearmanCorrCoeff[r, c] = 1.0
                if ENABLE_PEARSON_METHOD_1:
                    PearsonCorrCoeff[r, c] = 1.0
            else:
                ccol = HD5dat['data/vertices'][:,VARIABLE_FLUX_NAMES_idx[c]]
                pcc = 0.0
                if ENABLE_SPEARMAN_METHOD_1:
                    pcc = scipy.stats.spearmanr(rcol,  ccol)
                    SpearmanCorrCoeff[r, c] = pcc[0]
                    SpearmanCorrCoeff[c, r] = pcc[0]
                    print(r, c, pcc, 'spearman')
                if ENABLE_PEARSON_METHOD_1:
                    pcc = scipy.stats.pearsonr(rcol,  ccol)
                    PearsonCorrCoeff[r, c] = pcc[0]
                    PearsonCorrCoeff[c, r] = pcc[0]
                    print(r, c, pcc, 'pearson')
    print('PEARSONorSPEARMAN_METHOD_1 time:', float(time.time()-TIME_START)/60.0, 'minutes')

if ENABLE_SPEARMAN_METHOD_1:
    outputCorrCoef(SpearmanCorrCoeff, 'spearman', VARIABLE_FLUX_NAMES, model_name=model_name, work_dir=work_dir)
if ENABLE_PEARSON_METHOD_1:
    outputCorrCoef(PearsonCorrCoeff, 'pearson', VARIABLE_FLUX_NAMES, model_name=model_name, work_dir=work_dir)

HD5dat.close()




























